import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

# 设置复平面范围
x = np.linspace(-2, 2, 500)  # 实部范围
y = np.linspace(-2, 2, 500)  # 虚部范围
X, Y = np.meshgrid(x, y)
Z = X + 1j * Y  # 创建复数网格

# 计算平方根（主分支）
with np.errstate(invalid='ignore'):  # 忽略负实轴的警告
    W = np.sqrt(Z)

# 提取实部和虚部
real_part = np.real(W)
imag_part = np.imag(W)

# 创建3D图形
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 绘制曲面（使用虚部作为颜色）
surf = ax.plot_surface(
    X, Y, real_part, 
    facecolors=cm.jet((imag_part - imag_part.min()) / (imag_part.max() - imag_part.min())),
    rstride=2, cstride=2,  # 降低分辨率以提高性能
    antialiased=True
)

# 设置标签和标题
ax.set_xlabel('Re(z)', fontsize=12)
ax.set_ylabel('Im(z)', fontsize=12)
ax.set_zlabel('Re($\sqrt{z}$)', fontsize=12)
ax.set_title('$w = \sqrt{z}$: Height = Re(w), Color = Im(w)', fontsize=14)

# 添加颜色条（表示虚部）
mappable = cm.ScalarMappable(cmap=cm.jet)
mappable.set_array(imag_part)
cbar = fig.colorbar(mappable, ax=ax, shrink=0.7)
cbar.set_label('Im($\sqrt{z}$)', fontsize=12)

# 设置视角
ax.view_init(elev=30, azim=-60)  # 仰角30度，方位角-60度

# 保存高分辨率图像
plt.savefig('complex_sqrt.png', dpi=300, bbox_inches='tight')
plt.show()